import ollama
from openai import OpenAI
from rich.live import Live
from rich.console import Console
from rich import print
from rich.markdown import Markdown
from .config import read_config

console = Console()

def openai_chat(message: str, api: str, model: str) -> None:
    """### OpenAI chat"""
    env = read_config()
    api_key = env.get(f'{api}_api_key')
    base_url = env.get(f'{api}_base_url')
    try:
        client = OpenAI(api_key=api_key, base_url=base_url)
        response = client.chat.completions.create(
            model=model,
            messages=[{'role': 'user', 'content': message}],
            stream=True
        )
        full_content = ''
        with Live(console=console, auto_refresh=False) as live:
            for chunk in response:
                if chunk.choices and chunk.choices[0].delta:
                    content_chunk = chunk.choices[0].delta.content or ''
                    full_content += content_chunk
                    markdown = Markdown(markup=full_content, code_theme='material')
                    live.update(markdown, refresh=True)
            print()
    except Exception as e:
        print(f'[bold red]{str(e)}[/bold red]')

def ollama_chat(content: str, model: str) -> None:
    """### Ollama chat"""
    full_content = ''
    try:
        response = ollama.chat(
            model=model,
            messages=[{'role': 'user', 'content': content}],
            stream=True
        )
        with Live(console=console, auto_refresh=False) as live:
            for chunk in response:
                if chunk['message']['content']:
                    content_chunk = chunk['message']['content']
                    full_content += content_chunk
                    markdown = Markdown(markup=full_content, code_theme='material')
                    live.update(markdown, refresh=True)
            print()
    except Exception as e:
        print(f'[bold red]{str(e)}[/bold red]')
